import joblib
import numpy as np
import pandas as pd

# 加载模型和向量化器
model = joblib.load("text_classifier.pkl")
vectorizer = joblib.load("tfidf_vectorizer.pkl")

# 从训练数据中自动构建类别→非遗名称的映射
try:
    train_data = pd.read_csv("任务四词语.csv", encoding='utf-8-sig')
    category_to_name = dict(zip(train_data["类别"], train_data["非遗名称"]))
    print("自动加载的类别映射:", category_to_name)
except Exception as e:
    print("加载训练数据失败:", e)
    category_to_name = {}  # 如果失败，使用空字典

# 输入文本
input_text = "王明介绍制陶过程由挖陶、晒陶土、制坯、干燥等步骤组成"  # 可以替换成任意文本

# 文本向量化
X_test = vectorizer.transform([input_text])

# 预测
predicted_label = model.predict(X_test)[0]  # 预测类别
predicted_proba = model.predict_proba(X_test)[0]  # 各类别概率
confidence = np.max(predicted_proba)  # 最高置信度

# 获取非遗名称（优先从映射中获取，否则返回"未知非遗名称"）
predicted_name = category_to_name.get(predicted_label, "未知非遗名称")

# 输出结果
print("\n=== 预测结果 ===")
print(f"输入文本: {input_text}")
print(f"预测类别: {predicted_label}")
print(f"非遗名称: {predicted_name}")
print(f"置信度: {confidence:.4f}")

# 打印所有类别及其概率（调试用）
#print("\n所有类别概率:")
#for class_name, prob in zip(model.classes_, predicted_proba):
    #print(f"- {class_name}: {prob:.4f}")